//
// Created by wangwanlong on 2024/11/1.
//

#ifndef MKAHYPAR_PARTITIONER_H
#define MKAHYPAR_PARTITIONER_H

#include "mkahypar/utils/typedef.h"
#include "mkahypar/partition/coarsening/coarsener_base.h"
#include "mkahypar/partition/coarsening/ml_coarsener.h"
//#include "mkahypar/partition/initial_partitioning/random_initial_partitioner.h"
#include "mkahypar/partition/initial_partitioning/search0_initial_partitioner.h"
#include "mkahypar/partition/refinement/simple_refiner.h"
//负责整个划分流程
namespace mkahypar {
  class Partitioner {
  public:
    Partitioner(Hypergraph &hg, TargetGraph &topo) : _hg(hg), _topo(topo) {
      coarsener = std::make_unique<MLCoarsener<>>(hg,topo);
      initialPartitioner = std::make_unique<Search0InitialPartitioner>(hg, topo);
      refiner = std::make_unique<SimpleRefiner<NormalBlockPenalty, NormalEdgePenaltyPolicy, NormalHopPenaltyPolicy, AdvancedRandomWalkModelStopsSearch>>(
              hg, topo);
    }

    //设置一些动态参数
    void inline setupContext() {
      //-----------------设置缩合参数------------
      //允许缩合后的最大节点权重不超过单板的
      Context::getInstance().max_allowed_node_weight = _topo._fpga[0]._bound_res;
      Context::getInstance().max_allowed_node_weight /= Context::getInstance().proportion;
      LOG << "fpga coarsen bound:";
      Context::getInstance().max_allowed_node_weight.print();
      //缩合limit
      Context::getInstance().contraction_limit = 300;
    }

    //预处理
    void preprocess() {

    }

    void calDistance() {
      //提前计算出fpga之间的最短路
      _topo.calDistance();
    }

    //检查数据结构计算是否正确
    void checkDataStructure() {
      int num_partitions = _hg._k;
      auto &block = _topo._fpga;
      //检查block的资源是否正确
      std::vector<NodeWeight> res_used(num_partitions);
      for (const HypernodeID &hn: _hg.nodes()) {
        res_used[_hg.partID(hn)] += _hg.nodeWeight(hn);
      }
      for (int i = 0; i < num_partitions; ++i) {
        for (int j = 0; j < res_used.size(); ++j) {
          if (res_used[i][j] != block[i]._used_res[j]) {
            LOG << "type" << j << "not correct!";
            LOG << res_used[i][j] << block[i]._used_res[j];
            ERROREXIT("");
          }
        }
      }
      //检查超边的数据结构
      for (const HyperedgeID &he: _hg.edges()) {
        int edge_size = 0;
        auto source = _hg.hyperedge(he).source();
        std::vector<int> pin_in_part_size(num_partitions);
        std::vector<int> part_route_count(num_partitions);
        int64_t connect = 0;
        for (auto hn: _hg.pins(he)) {
          ++edge_size;
          pin_in_part_size[_hg.partID(hn)] += 1;
          connect |= 1LL << _hg.partID(hn);
          if (hn != source)
            part_route_count[_hg.partID(hn)] += _hg.tarRouteCount(he, hn);
        }
        //检查联通集
        if (connect != _hg.hyperedge(he).connectivity._n) {
          LOG << "connectivity  error!";
          LOG << he << connect << _hg.hyperedge(he).connectivity._n;
          ERROREXIT("");
        }
        //检查edge size
        if (edge_size != _hg.edgeSize(he)) {
          LOG << "edge size error!";
          LOG << he << edge_size << _hg.edgeSize(he);
          ERROREXIT("");
        }
        for (int id = 0; id < num_partitions; ++id) {
          //检查pin size
          if (pin_in_part_size[id] != _hg.pinCountInPart(he, id)) {
            LOG << "pin in part size error!";
            LOG << he << pin_in_part_size[id] << _hg.pinCountInPart(he, id);
            ERROREXIT("");
          }
          //检查part route count
          if (part_route_count[id] != _hg.routeCountInPart(he, id)) {
            LOG << "route count in part size error!";
            LOG << he << id << part_route_count[id] << _hg.routeCountInPart(he, id);
            ERROREXIT("");
          }
        }
      }
      //检查超点的数据结构
      for (const HypernodeID &hn: _hg.nodes()) {
        int num_incident_cuts = 0;
        for (const auto &he: _hg.hypernode(hn).incidentNets()) {
          if (_hg.connectivity(he) > 1)
            ++num_incident_cuts;
        }
        if (num_incident_cuts != _hg.hypernode(hn).num_incident_cut_hes) {
          LOG << "node cut error!";
          LOG << hn << num_incident_cuts << _hg.hypernode(hn).num_incident_cut_hes;
          ERROREXIT("error!");
        }
      }

      //检查互联量
      std::vector<int> line(num_partitions);
      for (const HyperedgeID &he: _hg.edges()) {
        auto he_weight = _hg.hyperedge(he).weight();
        if (_hg.connectivity(he) > 1) {
          for (int id = 0; id < num_partitions; ++id) {
            if(_hg.pinCountInPart(he,id) > 0)
              line[id] += he_weight;
          }
        }
      }
      for (int id = 0; id < num_partitions; ++id) {
        if(line[id] != _topo._fpga[id]._used_line){
          LOG<<"line error!";
          LOG<<id<<line[id]<<_topo._fpga[id]._used_line;
//          ERROREXIT("error!");
        }
      }

    }

    //打印信息
    void prinInfo() {
      //打印每个block的使用情况
      auto &num_partitions = _hg._k;
      std::vector<double> percents(num_partitions, 0);
      for (int i = 0; i < num_partitions; i++) {
        LOG << "block id: " << i;
        LOG << "line used:" << _topo._fpga[i]._used_line;
        LOG << "line bound:" << _topo._fpga[i]._bound_line;
        std::cout << "res used:";
        _topo._fpga[i]._used_res.print();
        std::cout << "res bound:";
        _topo._fpga[i]._bound_res.print();
        std::cout << "res remain:";
        (_topo._fpga[i]._bound_res - _topo._fpga[i]._used_res).print();

        // 找到资源占用率最高的一个资源，作为使用百分比
        std::pair<double, int> max_percent_type = _topo._fpga[i].getResUsedPercent();
        LOG << "the percent:" << max_percent_type.first * 100 << " the type:" << max_percent_type.second;
        percents[i] = max_percent_type.first;
      }

    }

    Metrics partition() {
//      LOG<<"start partiton:";
//      _hg.printGraph();
      //执行缩合
//      LOG << "before coarsen cost time:" << Timer::getInstance().getCostTime();
      coarsener->coarsen(Context::getInstance().contraction_limit);
//      LOG << "coarsen finish cost time:" << Timer::getInstance().getCostTime();

      //计算初始解
//      LOG << "before initialize cost time:" << Timer::getInstance().getCostTime();
      initialPartitioner->partition();
//      LOG << "initialize finish cost time:" << Timer::getInstance().getCostTime();

//      prinInfo();
//      _hg.printGraph();
//      checkDataStructure();//检查

      //执行还原
//      LOG << "before uncoarsen cost time:" << Timer::getInstance().getCostTime();
      Metrics best;
      coarsener->uncoarsen(*refiner,best);



//      checkDataStructure();//检查
//      check route count
      for(const HyperedgeID& he : _hg.edges()){
        HypernodeID source = _hg.hyperedge(he).source();
        for(const auto& tar : _hg.pins(he)){
          if(tar == source)
            continue;
          if(_hg.edgeTarRouteCount(he)[tar] != 1){
            ERROREXIT("error!");
          }
        }
      }

      return best;
    }



    std::unique_ptr<IInitialPartitioner> initialPartitioner;
    std::unique_ptr<ICoarsener> coarsener;
    std::unique_ptr<IRefiner> refiner;
    TargetGraph &_topo;
    Hypergraph &_hg;
  };
}
#endif //MKAHYPAR_PARTITIONER_H
